package com.googlecode.marrowboy.agent;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

public class HamcrestMonkeyPatcher implements ClassFileTransformer {

	public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] classfileBuffer) throws IllegalClassFormatException {

		if (className.equals("org/hamcrest/MatcherAssert")) {

			try {

				CtClass ct = ClassPool.getDefault().makeClass(new ByteArrayInputStream(classfileBuffer));

				for (CtMethod meth : ct.getDeclaredMethods()) {

					if ("assertThat".equals(meth.getName())) {

						if (meth.getParameterTypes().length == 2) {
							meth.insertBefore("com.googlecode.marrowboy.FakeMarrowBoyAssert.assertThat($1, $2);");
						}

						if (meth.getParameterTypes().length == 3) {
							meth.insertBefore("com.googlecode.marrowboy.FakeMarrowBoyAssert.assertThat($1, $2, $3);");
						}

					}

				}

				System.out.println("[INFO] MarrowBoy monkey-patch of org.hamcrest.MatcherAssert was successful!");
				
				return ct.toBytecode();

			} catch (Exception e) {
				e.printStackTrace();
				throw new RuntimeException(e);
			}

		}

		return null;
	}
}
